from scipy.constants import c
import torch
import numpy as np
import matplotlib.pyplot as plt

from BinNumCalZ import bin_num_cal_vectorized


def truncated_gauss(shape: list, loc: float = 0.0, scale: float = 1.0, a: float = -3.0, b: float = 3.0,
                    device='cpu') -> torch.Tensor:
    """
    使用torch库生成截断高斯分布
    :param shape: 形状
    :param loc: 均值
    :param scale: 标准差
    :param a: 下界
    :param b: 上界
    :param device: 'cpu'/'cuda' 根据测试，在生成的时候，第一次计算选cpu在5000×800的大小下比cuda快50倍，此后都是GPU快，但是生成坐标操作次数少，所以建议使用GPU，除非大规模调用此函数
    :return: torch Tensor
    """
    sqrt2 = torch.sqrt(torch.tensor(2.0, device=device))
    # 直接计算，让 PyTorch 处理类型转换
    phi_a = 0.5 * (1 + torch.erf((a - loc) / (scale * sqrt2)))
    phi_b = 0.5 * (1 + torch.erf((b - loc) / (scale * sqrt2)))
    u = torch.rand(shape, device=device) * (phi_b - phi_a) + phi_a
    x = loc + scale * sqrt2 * torch.erfinv(2 * u - 1)
    return torch.clamp(x, a, b)


def _resonator_coefficients(q) -> tuple:
    """
    返回纵向谐振子尾场的两个系数
    :param q: 品质因子
    :return: imag_factor, rot_factor sin项前面的系数, 旋转频率/谐振子频率即w_n / w_r
    """
    return 1 / np.sqrt(4 * q ** 2 - 1), np.sqrt(1 - 1 / (4 * q ** 2))


def resonator_wakel(omega_r: float, r: float, q: float, t: np.ndarray) -> np.ndarray:
    """
    生成谐振子尾场
    :param omega_r: 谐振频率
    :param r: 分路阻抗
    :param q: 品质因子
    :param t: 计算时间, 以粒子为参考, 跑在粒子前面的为负值, 后面的为正值, 需要按顺序排列好
    :return: np.ndarray
    """
    negatives = t[t < 0]
    wake_minus = np.zeros(negatives.shape)  # 相对论因果性
    t = t[t >= 0]
    imag_factor, rot_factor = _resonator_coefficients(q)
    constant = -omega_r * r / q  # 谐振子尾场前面的常数项系数
    decay = np.exp(-t * omega_r / 2 / q)  # 谐振子尾场衰减项
    rot_cos = np.cos(t * omega_r * rot_factor)  # cos项
    rot_sin = - imag_factor * np.sin(t * omega_r * rot_factor)  # sin项
    wake = constant * decay * (rot_cos + rot_sin)
    id0 = t == 0
    wake[id0] = wake[id0] / 2  # 设置粒子坐标处的尾场
    wake = np.array([*wake_minus, *wake])  # 拼接两段
    return wake


def total_short_wakel(t, *args) -> np.ndarray:
    """
    :param t: 统一的时间序列
    :param args: 传入的尾场数数组, 可以传入多个, 要求: 第一列为时间，第二列为尾场
    :return: 总的纵向尾场, 超出输入尾场范围的设置为0
    """
    wake = np.zeros(t.shape)
    for arg in args:
        wake += np.interp(t, arg[:, 0], arg[:, 1], left=0, right=0)
    return wake


def total_short_waket(t, *args) -> np.ndarray:
    """横向尾场加权求和"""


class Accelerator:
    def __repr__(self):
        return 'Accelerator'

    def __str__(self):
        max_key_length = max(len(key) for key in self.__dict__.keys())
        sorted_dict = dict(sorted(self.__dict__.items(), key=lambda item: item[0]), reverse=False)
        information = [f'{key:-<{max_key_length}}: {value:.2e}' for key, value in sorted_dict.items()]
        return '\n'.join(information)

    def __init__(self, parameters: dict):
        """
        初始化的时候会根据输入把储存环的一些基本属性计算好，除了输入的几个属性，还有：
        环相关的: 半径, 回旋周期, 回旋频率 radius, t0, f0
        腔压与能量相关的 : 自然同步相位 fais_nat
        辐射阻尼与量子激发相关的: 辐射阻尼率, 辐射阻尼系数, 量子激发率, 量子激发系数 lambda_rads, lambda_radz, radampcoef, quanexcoef
        还有两个作暂时不知道是什么的系数 kick_coef， ploss
        :param: parameters 应该包含self.__dict__ = parameters这行代码前的所有内容，属性名称为'键'，值为内容
        """
        # 环相关
        self.circumference = 0  # 周长/m
        self.alpha_c = 0  # 动量紧缩因子
        # 腔压与能量
        self.v_mc = 0  # 主腔腔压/V
        self.u0 = 0  # 能量损失/eV
        self.e0 = 0  # 电子能量/eV
        self.h = 0  # 谐波数
        # 束团相关
        self.i0 = 0  # 流强/A
        self.sigma_t0 = 0  # 束团归一化长度
        self.sigma_e0 = 0  # 束团归一化能散
        # 辐射阻尼与量子激发
        self.tau_s = 0  # 阻尼时间/s
        self.tau_z = 0  # 量子激发时间/s
        self.__dict__ = parameters  # 直接接收用户输入的系数
        # 环相关
        self.radius = self.circumference / 2 / np.pi  # 半径/m
        self.t0 = self.circumference / c  # 回旋周期/s^-1
        self.f0 = c / self.circumference  # 回旋频率/s
        self.tb = self.t0 / self.h  # 相邻bucket间距
        # 腔压与能量
        self.fais_nat = np.pi - np.asin(self.u0 / self.v_mc)  # 自然同步相位
        # 束团相关
        self.drift_coef = self.alpha_c * self.t0 * self.sigma_e0 / self.sigma_t0  # 单圈坐标迭代时, 动量引起的坐标变化即: q = drift_coef * p
        self.kick_coef = 1 / (self.e0 * self.sigma_e0)  # 用于求尾场kick的时候来省略一步乘法
        self.ploss = self.kick_coef * self.u0  # 单圈能量损失
        # 辐射阻尼与量子激发
        self.lambda_rads = self.t0 / self.tau_s  # 辐射阻尼率/圈
        self.lambda_radz = self.t0 / self.tau_z  # 量子激发率/圈
        self.radampcoef = 2 * self.lambda_rads  # 辐射阻尼系数
        self.quanexcoef = 2 * np.sqrt(self.lambda_radz)  # 量子激发系数


class Cavity:
    def __repr__(self):
        return 'Cavity'

    def __str__(self):
        max_key_length = max(len(key) for key in self.__dict__.keys())
        sorted_dict = dict(sorted(self.__dict__.items(), key=lambda item: item[0]), reverse=False)
        information = [f'{key:-<{max_key_length}}: {value}' for key, value in sorted_dict.items()]
        return '\n'.join(information)

    def __init__(self, v: float, f: float, q: float, r: float, detuning: float = 0):
        """
        创建一个腔, 除了输入的属性外，还会计算谐振频率(角频率和频率都有), 尾场的系数
        :param v: 腔压/V
        :param f: 工作频率/Hz
        :param q: 品质因子
        :param r: 分路阻抗/Ohm
        :param detuning: 失谐量/Hz
        """
        self.v = v
        self.f = f
        self.omega = self.f * 2 * np.pi
        self.q = q
        self.r = r
        self.detuning = detuning

        # 计算腔的谐振频率
        self.f_r = self._cavity_detuning(detuning)
        self.omega_r = self.f_r * 2 * np.pi

        # 腔的谐振子尾场的系数
        self.imag_factor, self.rot_factor = _resonator_coefficients(q)

    def _cavity_detuning(self, detuning):
        """返回腔失谐后的共振频率"""
        if detuning == 0:
            print('warning: 腔没有失谐')
        return self.f + detuning

    def short_wakel(self, t):
        """返回短程谐振子尾场， 长程尾场另作处理"""
        return resonator_wakel(self.omega_r, self.r, self.q, t)

    def pi_feedback(self):
        """添加PI反馈"""


class Beam:
    def __repr__(self):
        return 'Beam'

    def __str__(self):
        max_key_length = max(len(key) for key in self.__dict__.keys())
        sorted_dict = dict(sorted(self.__dict__.items(), key=lambda item: item[0]), reverse=False)
        information = [f'{key:-<{max_key_length}}: {value}' for key, value in sorted_dict.items()]
        return '\n'.join(information)

    def __init__(self, par_num, fill_pattern: np, q_range, p_range, device='cpu'):
        """
        创建一个粒子类，用于存放粒子坐标以及属性
        :param par_num: 切片数目
        :param fill_pattern: 填充模式, 输入的时候请转化为numpy数组
        :param q_range: 坐标范围
        :param p_range: 动量范围
        :param device: 传递给torch库的device参数, 用于生成CPU或GPU数组
        """
        self.fillrate = np.mean(fill_pattern)
        self.bin_num = int(self.fillrate * fill_pattern.shape[1])  # 计算非空bucket数目
        self.par_num = par_num
        self.device = device
        self.q = truncated_gauss(shape=[par_num, self.bin_num], loc=(q_range[0] + q_range[1]) / 2, scale=1,
                                 a=q_range[0], b=q_range[1], device=device)
        self.p = truncated_gauss(shape=[par_num, self.bin_num], loc=(p_range[0] + p_range[1]) / 2, scale=1,
                                 a=q_range[0], b=q_range[1], device=device)

    def drift_pass(self, ring: Accelerator):
        """老师代码里面先更新坐标再更新动量"""
        self.q = self.q - self.p * ring.drift_coef
        self.p = self.p - ring.ploss

    def cavity_kick(self, cavity: Cavity, ring: Accelerator):
        self.p = self.p + cavity.v * ring.kick_coef * torch.sin(-cavity.omega * ring.sigma_t0 * self.q + ring.fais_nat)

    def damping_excitation(self, ring: Accelerator):
        self.p = self.p + ring.radampcoef * self.p + ring.quanexcoef * torch.rand([self.par_num, self.bin_num],device=self.device)
